# This test script performs sharpening of the given input image.
# This is a good example of a fairly heavy operation that uses a
# large input and output.

from __future__ import with_statement
import sys
import scavenger
from time import sleep, time
import ImageFile

# Check whether a number of test iterations is given.
iterations = 1
if '-i' in sys.argv:
    iterations = int(sys.argv[sys.argv.index('-i') + 1])

# Check whether a scheduler has been specified.
scheduler = 'basic'
if '-s' in sys.argv:
    scheduler = sys.argv[sys.argv.index('-s') + 1]

# Check for the depth argument.
depth = 50
if '-d' in sys.argv:
    depth = int(sys.argv[sys.argv.index('-d') + 1])

# Sleep for a little while to make sure that we discover
# the available surrogates.
sleep(2.0)

width = 640
height = 480
start_coord = (-1.5+-1j) 
end_coord = (0.5+1j)

@scavenger.scavenge
def mandelbrot(start, end, width, height, iterations):
    data = []
    # Create the fractal.
    hstep = (end.real-start.real)/width
    vstep = (end.imag-start.imag)/height
    cstep = 255.0/iterations
    black = chr(255)
    for i in range (0, height):
        current_vstep = start.imag + vstep*i
        current_hstep = 0.0
        for j in range (0, width):
            z = (0+0j)
            c = complex(start.real + current_hstep, current_vstep)
            current_hstep += hstep
            color = 0.0
            k = 0
            while k < iterations:
                z *= z
                z += c
                if abs(z) > 2:
                    # We have reached the limit. Write a pixel of the current color.
                    outcol = int(color)
                    data.append(chr(outcol))
                    break
                color += cstep
                k += 1
            # If nothing was written in the loop we write a black pixel.
            if k == iterations:
                data.append(black)
    return ''.join(data)

@scavenger.dcscavenge('#2*#3', 200)
def mandelbrot_dc(start, end, width, height, iterations):
    data = []
    # Create the fractal.
    hstep = (end.real-start.real)/width
    vstep = (end.imag-start.imag)/height
    cstep = 255.0/iterations
    black = chr(255)
    for i in range (0, height):
        current_vstep = start.imag + vstep*i
        current_hstep = 0.0
        for j in range (0, width):
            z = (0+0j)
            c = complex(start.real + current_hstep, current_vstep)
            current_hstep += hstep
            color = 0.0
            k = 0
            while k < iterations:
                z *= z
                z += c
                if abs(z) > 2:
                    # We have reached the limit. Write a pixel of the current color.
                    outcol = int(color)
                    data.append(chr(outcol))
                    break
                color += cstep
                k += 1
            # If nothing was written in the loop we write a black pixel.
            if k == iterations:
                data.append(black)
    return ''.join(data)

@scavenger.gprofilescavenge('#2*#3')
def mandelbrot_gprofile(start, end, width, height, iterations):
    data = []
    # Create the fractal.
    hstep = (end.real-start.real)/width
    vstep = (end.imag-start.imag)/height
    cstep = 255.0/iterations
    black = chr(255)
    for i in range (0, height):
        current_vstep = start.imag + vstep*i
        current_hstep = 0.0
        for j in range (0, width):
            z = (0+0j)
            c = complex(start.real + current_hstep, current_vstep)
            current_hstep += hstep
            color = 0.0
            k = 0
            while k < iterations:
                z *= z
                z += c
                if abs(z) > 2:
                    # We have reached the limit. Write a pixel of the current color.
                    outcol = int(color)
                    data.append(chr(outcol))
                    break
                color += cstep
                k += 1
            # If nothing was written in the loop we write a black pixel.
            if k == iterations:
                data.append(black)
    return ''.join(data)
    
# Choose the function to call - based on the scheduler chosen.
function = None
if scheduler == 'basic':
    function = mandelbrot
elif scheduler == 'dc':
    function = mandelbrot_dc
elif scheduler == 'gprofile':
    function = mandelbrot_gprofile
elif scheduler == 'decomp':
    function = mandelbrot_decomp
else:
    raise Exception('Unknown scheduler.')

# Run the test!
for x in range(0, iterations):
    start = time()
    try:
        result = function(start_coord, end_coord, width, height, depth)
    except:
        print '-1',
    stop = time()
    print stop - start

# Write the result for possible inspection.
#p = ImageFile.Parser()
#p.feed("P6 %u %u 255\n"%(width,height))
#for x in result:
#    p.feed('%s%s%s'%(x,x,x))
#im = p.close()
#im.save("result.jpg")

scavenger.shutdown()
